import numpy as np
import cv2
import json
import argparse
import os
from pathlib import Path
import shutil
from tqdm import tqdm
import json

def labelme2mask_single_img(img_path, labelme_json_path, class_info):
    '''
    输入原始图像路径和labelme标注路径, 输出 mask
    '''
    img_bgr = cv2.imread(img_path)
    img_mask = np.zeros(img_bgr.shape[:2]) # 创建空白图像 0-背景
    with open(labelme_json_path, 'r', encoding='utf-8') as f:
        labelme = json.load(f)

    for one_class in class_info: # 按顺序遍历每一个类别
        for each in labelme['shapes']: # 遍历所有标注，找到属于当前类别的标注
            if each['label'] == one_class['label'] and each['shape_type'] == one_class['type']:
                if one_class['type'] == 'polygon': # polygon 多段线标注
                    # 获取点的坐标
                    points = [np.array(each['points'], dtype=np.int32).reshape((-1, 1, 2))]
                    # 在空白图上画 mask（闭合区域）
                    img_mask = cv2.fillPoly(img_mask, points, color=one_class['color'])
                     # 在多边形边缘画线
                    img_mask = cv2.polylines(img_mask, points, isClosed=True, color=255, thickness=3)
                elif one_class['type'] == 'line' or one_class['type'] == 'linestrip': # line 或者 linestrip 线段标注
                    # 获取点的坐标
                    points = [np.array(each['points'], dtype=np.int32).reshape((-1, 1, 2))]
                    # 在空白图上画 mask（非闭合区域）
                    img_mask = cv2.polylines(img_mask, points, isClosed=False, color=255, thickness=one_class['thickness']+3)
                    img_mask = cv2.polylines(img_mask, points, isClosed=False, color=one_class['color'], thickness=one_class['thickness'])
                else:
                    print('未知标注类型', one_class['shape_type'])
    return img_mask


def Parse_Arguments():
    parser = argparse.ArgumentParser(description="")
    parser.add_argument("--input_dir", type=str, default="", help="Input annotated directory")
    parser.add_argument("--output_dir", type=str, default="", help="Output dataset directory")
    parser.add_argument("--label_path", type=str, default="myutil/make_dataset/label.json", help="label json file path")
    return parser.parse_args()
def main():
    args = Parse_Arguments()
    # 0-背景，从 1 开始
    class_info = {}
    if Path(args.label_path).exists():
        with open(args.label_path, "r", encoding="utf-8") as json_file:
            class_info = json.load(json_file)
    else:
        print(f"label 文件不存在")
    if os.path.exists(args.output_dir):
        print("目标目录存在，清空目标目录", args.output_dir)
        shutil.rmtree(args.output_dir)

    os.makedirs(args.output_dir)
    os.makedirs(os.path.join(args.output_dir, "JPEGImages"))
    os.makedirs(os.path.join(args.output_dir, "SegmentationClass"))
    image_paths = list(Path(args.input_dir).glob("*.jpg"))
    print("Creating dataset:", args.output_dir)

    for img_path in tqdm(image_paths):
        try:
            labelme_json_path = img_path.with_suffix(".json")
            img_mask = labelme2mask_single_img(str(img_path), str(labelme_json_path), class_info)
            mask_path = os.path.join(args.output_dir, "SegmentationClass",f"{img_path.with_suffix('.png').name}")
            cv2.imwrite(mask_path, img_mask)
            shutil.copy(img_path, os.path.join(args.output_dir, "JPEGImages"))
        except Exception as E:
            print(img_path, '转换失败', E)

if __name__ == "__main__":
    main()
